from apscheduler.schedulers.asyncio import AsyncIOScheduler
from apscheduler.triggers.interval import IntervalTrigger
from fastapi import Request, Depends
from functools import wraps
from typing import Callable, Any

class TaskScheduler:
    def __init__(self):
        self.scheduler = AsyncIOScheduler()
        self._route_tasks = []

    def task(self, trigger: str = "interval", **trigger_args):
        """
        支持双重功能的路由/任务装饰器
        """
        def decorator(func: Callable) -> Callable:
            @wraps(func)
            async def wrapper(
                    request: Request = None,
                    params: dict = None,
                    query_db: Any = Depends(None),
                    *args, **kwargs
            ):
                # 作为路由时的处理
                if request is not None:
                    return await func(request, params, query_db, *args, **kwargs)
                # 作为定时任务时的处理
                else:
                    return await func(None, params, None, *args, **kwargs)

            # 注册为定时任务
            if trigger == "interval":
                self._route_tasks.append({
                    "func": wrapper,
                    "trigger": IntervalTrigger(**trigger_args)
                })
            return wrapper
        return decorator

    def start(self):
        """启动所有路由任务"""
        for job in self._route_tasks:
            self.scheduler.add_job(
                job["func"],
                trigger=job["trigger"],
                name=job["func"].__name__
            )
        self.scheduler.start()

scheduler = TaskScheduler()